### Functions for calculating cross elas matrix ###
### MCES version
rich.subs <- function(dat,eqbm.data,market) { #dat = output from DGP.exog()
  #set verbose = TRUE if you want to see the Iter and Residual  for the FPITER.
   #
  source("eqbm_subfunctions.R",local=TRUE)
  # new functions for rich subs
  wtd.var.prm <- function(prm) {
    wt <- y.h/Y
    mu <- colSums(prm*wt)  
    colSums(wt*(t(t(prm)-mu))^2) 
  }
  #
  # crosselas : based on build.DeltaX, a matrix of cross price elasticities, changing price of m (row), impact on j(cols)
  crosselas.mat <- function(prm) {
    Delta <- matrix(0,nrow=n.models,ncol=n.models)
    for(m in 1:n.models) {
      for(j in 1:n.models) {
        Delta[m,j] = sum(alpha.h*prm[,m]*omega(prm,j))  # epsilon_{jm} * Omega_{jm} in theory
      }
    }
    return(Delta)
  }
  # start of eqbm computation #  
  df.m <- data.frame(subset(dat$models,n==market))
  i <- df.m$i
  if(U0==0) xvars <- paste0("x",1:n.x) else xvars <- paste0("x",0:n.x)
  x <- as.matrix(df.m[,xvars]) # if df.m is only a data.frame
  colnames(x) <- xvars  # no effect for n.x>1, but names it x1 instead of just x
  xi <- df.m$xi
  mc <- df.m$mc * (df.m$tau + df.m$tar.change) # apply trade costs and name it mc, for consistency with subordinate functions
  #
  df <- subset(dat$households,n==market)
  if(U0==0) bvec <- paste0("b",1:n.x,".h") else bvec <- paste0("b",0:n.x,".h")
  b.h <- as.matrix(df[,bvec]) 
  colnames(b.h) <- bvec  # no effect for n.x>1, but names it b1.h instead of just b.h
  alpha.h <- df$alpha.h
  y.h <- df$y.h
  Y <- sum(df$y.h)
  p <- eqbm.data$p # model level price vector
  wvp <-  wtd.var.prm(prob.mat(p)) # model level market share vector
  pelas.m <- ownelas(prm=prob.mat(p)) # model level own price elasticities vector
 # model-level variables
  m.eqvar <- data.table(m=eqbm.data$m,f_m=eqbm.data$f,s_m=eqbm.data$s,wvp_m=wvp,p_m=p,own_elas_m=pelas.m) # model level output for model m
  # m by j
  cross.elas <- as.data.table(crossElasMCES(MM=co.own.mat, HH=prob.mat(p), A=alpha.h, S=eqbm.data$s*Y, y=y.h))
  #
  #reshape and merge
  cross.elas[, m := 1:.N]
  cross.elas.long <- melt(cross.elas,measure.vars=patterns("^V"),value.name="cross_elas",id="m")
  cross.elas.long[,j := as.numeric(sub("V","",variable))]
  cross.elas.long[,variable := NULL]
  #merge ownelas
  cross.elas.long <- merge(cross.elas.long,m.eqvar,by="m",all.x=TRUE)
  #
  return(cross.elas.long)
}
#Monte function
monte.func.rich.subs <- function(repi,settings,v=1,market=1) { # monte carlo stage
set.seed(seedn)
#data
  data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
#eqbm
  eqbm.pre.mm  <- as.data.table(DGP.endog(parallel=F,v=v,tar.change=0,data.exog = data.exog.mm))
  eqbm.pre.mm.with.ces <- as.data.table(estapprox.mm(eqbm.pre.mm)) #add ces approximation results
  EQ <- eqbm.pre.mm.with.ces[n==market,] #keep only one destination
  EQ <- as.data.frame(EQ)
#MD
  if (v==3) {
    MD <- MahaDist(data=EQ,features = c(paste0("x",1:n.x),"p"))
  } else MD <- MahaDist(data=EQ,features = c(paste0("x",1:n.x)))
#XE  
  XE <- rich.subs(dat=data.exog.mm,eqbm.data =EQ,market=market) 
#Merge
  my.elas <- merge(XE,MD,by=c("m","j"))
#   
  return(data.frame(repi=repi,v=v,eta=EQ$eta[1] ,my.elas))
}
#Doing multiple monte carlo reps for a single market
monte.rep.rich.subs <- function(parallel=TRUE,settings,v=1,market=1){ #get equilibrium market shares and prices for each market and append them
  if(parallel==T)  foreach(i=1:n.reps, .combine = 'rbind')  %dorng%  monte.func.rich.subs(i,settings,v,market) else foreach(i=1:n.reps, .combine = 'rbind')  %do%  monte.func.rich.subs(i,settings,v,market) 
}
#Function for running graphs. takes setting vars
graph.richsubs <- function(settings.chosen,v.chosen=1, market.chosen=1){
  graph.dt <- as.data.table(monte.func.rich.subs(1,settings=settings.chosen,v=v.chosen,market=market.chosen))
  graph.dt[,v:=v.chosen]
  graph.dt[,eta.v:=exp(settings.chosen$mu.a[v.chosen])]
  return(graph.dt)
}

graph.richsubs.X <- function(nx){
  n.x <<- nx
  graph.dt <-  graph.richsubs(settings.chosen=settings.var,v=2,market.chosen=1)
  return(data.table(nx,graph.dt))
}

###MLOG VERSION###

rich.subs.mlog <- function(dat,eqbm.data,market) { #dat = output from DGP.exog()
  #set verbose = TRUE if you want to see the Iter and Residual  for the FPITER.
  #
  source("eqbm_MLOG_subfunctions.R",local=TRUE)
  # 
  #household variance in market share of m
  wtd.var.prm <- function(prm) {
    wt <- rep(1/n.hh,n.hh)
    mu <- colSums(prm*wt)  
    colSums(wt*(t(t(prm)-mu))^2) 
  }
  #
  omega <- function(prm,m) {
    (prm[,m])/sum(prm[,m])  
  }
  # crosselas : based on build.DeltaX, a matrix of cross price elasticities, changing price of m (row), impact on j(cols)
  crosselas.mat <- function(prm,price) {
    # 
    Delta <- matrix(0,nrow=n.models,ncol=n.models)
    for(m in 1:n.models) {
      for(j in 1:n.models) {
        Delta[m,j] = sum(alpha.h*prm[,m]*omega(prm,j))  # epsilon_{jm} * Omega_{jm} in theory
      }
    }
    return(price*Delta) # multiplies the matrix m rows j columns, by prices of m
  }
  #
  # start of eqbm computation #  
  df.m <- data.frame(subset(dat$models,n==market))
  i <- df.m$i
  if(U0==0) xvars <- paste0("x",1:n.x) else xvars <- paste0("x",0:n.x)
  x <- as.matrix(df.m[,xvars]) # if df.m is only a data.frame
  colnames(x) <- xvars  # no effect for n.x>1, but names it x1 instead of just x
  xi <- df.m$xi
  mc <- df.m$mc * (df.m$tau + df.m$tar.change) # apply trade costs and name it mc, for consistency with subordinate functions
  #
  df <- subset(dat$households,n==market)
  if(U0==0) bvec <- paste0("b",1:n.x,".h") else bvec <- paste0("b",0:n.x,".h")
  b.h <- as.matrix(df[,bvec]) 
  colnames(b.h) <- bvec  # no effect for n.x>1, but names it b1.h instead of just b.h
  alpha.h <- df$alpha.h
  y.h <- df$y.h
  Y <- sum(df$y.h)
  p <- eqbm.data$p # model level price vector
  wvp <-  wtd.var.prm(prob.mat(p)) # model level market share vector
  pelas.m <- ownelas(prm=prob.mat(p),share=eqbm.data$s,price=p) # model level own price elasticities vector
  # model-level variables
  m.eqvar <- data.table(m=eqbm.data$m,f_m=eqbm.data$f,s_m=eqbm.data$s,wvp_m=wvp,p_m=p,own_elas_m=pelas.m) # model level output for model m
  # m by j
  cross.elas <- as.data.table(crossElasMLOG(MM=co.own.mat, HH=prob.mat(p), A=alpha.h, S=eqbm.data$s, P=p))
  #
  #reshape and merge
  cross.elas[, m := 1:.N]
  cross.elas.long <- melt(cross.elas,measure.vars=patterns("^V"),value.name="cross_elas",id="m")
  cross.elas.long[,j := as.numeric(sub("V","",variable))]
  cross.elas.long[,variable := NULL]
  #merge ownelas
  cross.elas.long <- merge(cross.elas.long,m.eqvar,by="m",all.x=TRUE)
  #
  return(cross.elas.long)
}


